import tensorflow as tf
import numpy as np
from PIL import Image, ImageDraw
import os

# 加载 MNIST 数据集
mnist = tf.keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# 合并训练集和测试集
all_images = np.concatenate((train_images, test_images), axis=0)
all_labels = np.concatenate((train_labels, test_labels), axis=0)

# 定义一个函数将数字图像组合成多位数图像
def combine_digits(digits, all_images, all_labels, target_size=(128, 128)):
    digit_images = []
    for digit in digits:
        digit_indices = np.where(all_labels == digit)[0]
        digit_index = np.random.choice(digit_indices)
        digit_image = all_images[digit_index]
        # 反转图像，实现白底黑字
        digit_image = 255 - digit_image
        digit_image = Image.fromarray(digit_image)

        # 如果数字是 6 或 9，添加下划线
        if digit in [6, 9]:
            draw = ImageDraw.Draw(digit_image)
            width, height = digit_image.size
            # 增加下划线宽度并往上移动一点
            line_y = height - 3  # 往上移动 3 个像素，可以根据需要调整
            line_width = 2  # 下划线宽度为 2 个像素，可以根据需要调整
            draw.line((0, line_y, width, line_y), fill=0, width=line_width)

        digit_images.append(digit_image)

    # 组合数字图像
    combined_width = sum([img.width for img in digit_images])
    combined_image = Image.new('L', (combined_width, max([img.height for img in digit_images])))
    x_offset = 0
    for img in digit_images:
        combined_image.paste(img, (x_offset, 0))
        x_offset += img.width

    # 调整大小到目标大小
    combined_image = combined_image.resize(target_size, Image.BICUBIC)

    # 转换为三色图
    combined_image = combined_image.convert('RGB')

    return combined_image

# 每个类别生成的图片数量
num_images_per_class = 200

base_dir = './combined_images'
os.makedirs(base_dir, exist_ok=True)

for i in range(100):  # 遍历 0 - 99
    digits = [int(d) for d in str(i)]
    # 创建每个数字的类别目录
    digit_dir = os.path.join(base_dir, str(i))
    os.makedirs(digit_dir, exist_ok=True)
    for j in range(num_images_per_class):
        image = combine_digits(digits, all_images, all_labels)
        save_path = os.path.join(digit_dir, f'{i}_{j}.jpg')
        image.save(save_path, 'JPEG')